﻿
using System;
using System.IO;
using System.Runtime.InteropServices;
using WindowsAPI.Service;

namespace WindowsAPI.PE
{
    public class PEFile
    {
        /// <summary>
        /// Reads in a block from a file and converts it to the struct
        /// type specified by the template parameter
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="reader"></param>
        /// <returns></returns>
        protected static T FromBinaryReader<T>(BinaryReader reader)
        {
            // Read in a byte array
            var bytes = reader.ReadBytes(Marshal.SizeOf(typeof(T)));

            // Pin the managed memory while, copy it out the data, then unpin it
            T theStructure;
            using(var handle = new UsingObject<GCHandle>(GCHandle.Alloc(bytes, GCHandleType.Pinned), h => h.Free()))
                theStructure = (T)Marshal.PtrToStructure(handle.Object.AddrOfPinnedObject(), typeof(T));
            //var handle = GCHandle.Alloc(bytes, GCHandleType.Pinned);
            //var theStructure = (T)Marshal.PtrToStructure(handle.AddrOfPinnedObject(), typeof(T));
            //handle.Free();

            return theStructure;
        }

        private readonly IMAGE_DOS_HEADER f_DosHeader;
        private readonly IMAGE_FILE_HEADER f_FileHeader;
        private readonly IMAGE_OPTIONAL_HEADER32 f_OptionalHeader32;
        private readonly IMAGE_OPTIONAL_HEADER64 f_OptionalHeader64;
        private readonly IMAGE_SECTION_HEADER[] f_ImageSectionHeaders;
        private readonly Data_Directory[] f_DataDirectories;

        /// <summary>Gets if the file header is 32 bit or not</summary>
        public bool Is32BitHeader
        {
            get
            {
                return (f_FileHeader.Characteristics & ImageCharacteristicsType.IMAGE_FILE_32BIT_MACHINE)
                    == ImageCharacteristicsType.IMAGE_FILE_32BIT_MACHINE;
            }
        }

        public IMAGE_DOS_HEADER DOSHeader { get { return f_DosHeader; } }


        /// <summary>Gets the file header</summary>
        public IMAGE_FILE_HEADER FileHeader { get { return f_FileHeader; } }

        /// <summary>Gets the optional header</summary>
        public IMAGE_OPTIONAL_HEADER32 OptionalHeader32 { get { return f_OptionalHeader32; } }

        /// <summary>
        /// Gets the optional header
        /// </summary>
        public IMAGE_OPTIONAL_HEADER64 OptionalHeader64 { get { return f_OptionalHeader64; } }

        public IMAGE_SECTION_HEADER[] ImageSectionHeaders { get { return f_ImageSectionHeaders; } }

        /// <summary>Gets the timestamp from the file header</summary>
        public DateTime TimeStamp
        {
            get
            {
                // Timestamp is a date offset from 1970
                var returnValue = new DateTime(1970, 1, 1, 0, 0, 0);

                // Add in the number of seconds since 1970/1/1
                returnValue = returnValue.AddSeconds(f_FileHeader.TimeDateStamp);
                // Adjust to local timezone
                returnValue += TimeZone.CurrentTimeZone.GetUtcOffset(returnValue);

                return returnValue;
            }
        }

        public PEFile(string FileName)
        {
            using(var stream = new FileStream(FileName, FileMode.Open, FileAccess.Read))
            {
                var reader = new BinaryReader(stream);
                f_DosHeader = FromBinaryReader<IMAGE_DOS_HEADER>(reader);

                // Add 4 bytes to the offset
                stream.Seek(f_DosHeader.lfanew, SeekOrigin.Begin);

                var lv_NtHeadersSignature = reader.ReadUInt32();
                if(lv_NtHeadersSignature != 17744)
                    throw new FormatException("Не найдена сигнатура IMAGE_NT_HEADERS");

                f_FileHeader = FromBinaryReader<IMAGE_FILE_HEADER>(reader);
                uint NumberOfRvaAndSizes;
                if(Is32BitHeader)
                {
                    f_OptionalHeader32 = FromBinaryReader<IMAGE_OPTIONAL_HEADER32>(reader);
                    NumberOfRvaAndSizes = f_OptionalHeader32.NumberOfRvaAndSizes;
                }
                else
                {
                    f_OptionalHeader64 = FromBinaryReader<IMAGE_OPTIONAL_HEADER64>(reader);
                    NumberOfRvaAndSizes = f_OptionalHeader64.NumberOfRvaAndSizes;
                }

                f_DataDirectories = new Data_Directory[NumberOfRvaAndSizes];
                for(var i = 0; i < NumberOfRvaAndSizes; i++)
                    f_DataDirectories[i] = FromBinaryReader<Data_Directory>(reader);

                var sections_count = f_FileHeader.NumberOfSections;
                f_ImageSectionHeaders = new IMAGE_SECTION_HEADER[sections_count];
                for(var i = 0; i < sections_count; i++)
                    f_ImageSectionHeaders[i] = FromBinaryReader<IMAGE_SECTION_HEADER>(reader);
            }
        }

        /// <summary>Gets the header of the .NET assembly that called this function</summary>
        /// <returns></returns>
        public static PEFile GetCallingAssemblyHeader()
        {
            // Get the path to the calling assembly, which is the path to the
            // DLL or EXE that we want the time of
            var file = System.Reflection.Assembly.GetCallingAssembly().Location;

            // Get and return the timestamp
            return new PEFile(file);
        }

        /// <summary>Gets the header of the .NET assembly that called this function</summary>
        /// <returns></returns>
        public static PEFile GetAssemblyHeader()
        {
            // Get the path to the calling assembly, which is the path to the
            // DLL or EXE that we want the time of
            var file = System.Reflection.Assembly.GetAssembly(typeof(PEFile)).Location;

            // Get and return the timestamp
            return new PEFile(file);
        }
    }
}
